package com.jin.webvideo.base.dc;

import com.jin.webvideo.base.dc.safe.whitelist.Whitelist;
import com.jin.webvideo.base.dc.safe.whitelist.WhitelistManager;
import com.jin.webvideo.base.util.CollectionsUtil;
import javassist.bytecode.ClassFile;
import javassist.bytecode.ConstPool;

import java.io.DataInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.*;

/**
 * 校验工具
 *
 * @author biao532
 * @date 2023/12/02 20:07
 **/
public class Verifier {
    /**
     * 使用到的类名 白名单校验
     * @param inputStream 字节码输入流
     * @throws IOException
     */
    public static void whiteList(InputStream inputStream, String className) throws IOException {
        ClassFile cf = parser(inputStream);
        // 校验包名是否允许
        if (!pkgName(cf)) {
            throw new IOException("该包禁止被使用");
        }
        Set<String> s = getDependencies(cf);
        List<String> forbidClass = new ArrayList<>(0);
        List<Whitelist<String>> list = WhitelistManager.singleton().get(WhitelistManager.BASE_CLASS, packageName(className));
        // 校验依赖是否被允许
        for (String name : s) {
            Whitelist<String> exit = CollectionsUtil.first(list, (l) -> l.contains(name));
            if (exit == null) {
                forbidClass.add(name);
            }
        }
        if (forbidClass.isEmpty()) {
            return ;
        }
        throw new IOException("代码中存在的以下类被禁止使用：" + forbidClass);
    }

    public static boolean pkgName(ClassFile cf) throws IOException {
        Set<String> allow = WhitelistManager.singleton().definedPkg();
        return allow.contains(packageName(cf.getName()));
    }

    private static String packageName(String line) throws IOException {
        int idx = line.lastIndexOf(".");
        if (idx == -1) {
            return line;
        }
        if (idx == 0) {
            throw new IOException("不合法的包名");
        }
        return line.substring(0, idx);
    }

    /**
     * 获取依赖的类
     * @param cf
     * @return
     * @throws IOException
     */
    public static Set<String> getDependencies(ClassFile cf) throws IOException {
        //java.lang.Object, java.lang.String, java.lang.Boolean, java.lang.Byte, java.lang.Character, java.lang.Short, java.lang.Integer, java.lang.Float, java.lang.Long, java.lang.Double
        //java.util.Arrays, java.util.Collections,
        ConstPool constPool = cf.getConstPool();
        Set<String> set = new HashSet<>();
        for (int ix = 1, size = constPool.getSize(); ix < size; ix++) {
            // 类的全限定名
            if (constPool.getTag(ix) == ConstPool.CONST_Class) {
                set.add(constPool.getClassInfo(ix));
                continue;
            }
            // 字段或者方法的名称和描述
            if (constPool.getTag(ix) == ConstPool.CONST_NameAndType) {
                int descriptorIndex = constPool.getNameAndTypeDescriptor(ix);
                String desc = constPool.getUtf8Info(descriptorIndex);

                // 截取L...;中的...
                int p = -1;
                while ((p = desc.indexOf("L", p)) != -1 && ++p < desc.length()) {
                    int np = desc.indexOf(';', p);
                    set.add(desc.substring(p, np).replace('/', '.'));
                    p = np;
                }
            }
        }
        return set;
    }

    private static ClassFile parser(InputStream inputStream) throws IOException {
        return new ClassFile(new DataInputStream(inputStream));
    }
}
